import pysam

class AlignStr:
    '''简化的比对结果类，可以直接输出对齐后的比对结果，\n
    index以0开始且前闭后开，但输出是以1开始，且前闭后闭'''
    def __init__(self, ref_align_str:str, query_align_str:str, query_quality:list[int]=None, ref_start_index:int=0, query_start_index:int=0) -> None:
        self.ref_seq = ref_align_str
        self.query_seq = query_align_str
        if(query_quality is None):
            self.query_quality = [40]*len(self.query_seq)
        else:
            self.query_quality = query_quality
        self.ref_start_index = ref_start_index
        self.ref_end_index = self.ref_start_index + len(ref_align_str) - ref_align_str.count('-')
        self.query_start_index = query_start_index
        self.query_end_index = self.query_start_index + len(query_align_str) - query_align_str.count('-')
        
    def getMutationStat(self):
        insert_pos_list:list[int] = list()
        delete_pos_list:list[int] = list()
        replace_pos_list:list[int] = list()
        pos_mutation_hash = {'insert':insert_pos_list,'delete':delete_pos_list,'replace':replace_pos_list}
        p, i = [0, self.ref_start_index]
        for p in range(len(self.ref_seq)):
            ref_base = self.ref_seq[p]
            query_base = self.query_seq[p]

            if(ref_base == query_base):
                i += 1
            elif(ref_base == '-'):
                ######插入
                insert_pos_list.append(i-1)
                # insert_base_list.append(query_base)
            elif(query_base == '-'):
                ######缺失
                delete_pos_list.append(i)
                i += 1
            else:
                ######错配
                if(query_base != 'N'):
                    replace_pos_list.append(i) 
                i += 1

        return pos_mutation_hash     

    def reference_sub(self, start_index, end_index, ignore_gap=True):
        if(start_index >= end_index):
            raise IndexError()
        if(ignore_gap and (start_index < self.ref_start_index or end_index > self.ref_end_index)):
            raise IndexError()
        
        # ref_seq, query_seq = ['', '']
        if(ignore_gap):
            match_index_list = [i for i,v in enumerate(self.ref_seq) if(v != '-')]
            match_index_list.append(match_index_list[-1] + 1)
            start_index = match_index_list[start_index-self.ref_start_index]
            end_index = match_index_list[end_index-self.ref_start_index]
        
        ref_seq = self.ref_seq[start_index:end_index]
        query_seq = self.query_seq[start_index:end_index]
            
        return AlignStr(ref_seq, query_seq, 
                        self.ref_start_index+start_index-self.ref_seq[:start_index].count('-'),
                        self.query_start_index+start_index-self.query_seq[:start_index].count('-')
                        )
    
    def __str__(self) -> str:
        ref_str = f"{self.ref_start_index+1}\t{self.ref_seq}\t{self.ref_end_index}"
        query_str = f"{self.query_start_index+1}\t{self.query_seq}\t{self.query_end_index}"
        return ref_str+"\n"+query_str
    
    def __len__(self):
        return len(self.ref_seq)


def getAlignBounderIndex(align:pysam.AlignedSegment, pr1_index, pr2_index):
    '''
    [pr1_index, pr2_index)
    '''
    align_pairs = align.get_aligned_pairs()
    i, j = [0, len(align_pairs) - 1]
    while(i<len(align_pairs) and (align_pairs[i][1] is None or align_pairs[i][1]<pr1_index)):
        i += 1

    if(align_pairs[j][1] is not None and align_pairs[j][1] < pr2_index):
        j += 1
    else:
        while(j>=0 and (align_pairs[j][1] is None or align_pairs[j][1]>pr2_index)):
            j -= 1
    return i, j

def alignToStr(align:pysam.AlignedSegment, refSeq:str=None, startIndex:int=0, endIndex:int=None)->AlignStr:
    ref_align_str_list:list[str]=[]
    query_align_str_list:list[str]=[]
    query_quality_list:list[str]=[]
    query_sequence = align.query_sequence
    if(refSeq is None):
        refSeq = align.get_reference_sequence()
        refSeq = '-'*align.reference_start + refSeq

    pairs = align.get_aligned_pairs()
    if(endIndex is None):
         endIndex = len(pairs)

    ref_start_index = None
    query_start_index = None
    for i in range(startIndex, endIndex):
        pair=pairs[i]###第i个匹配结果
        #######获取匹配的碱基，若出现None则用'-'代替
        ref_align_str_list.append('-' if pair[1] is None else refSeq[pair[1]])
        query_align_str_list.append('-' if pair[0] is None else query_sequence[pair[0]])
        query_quality_list.append(0 if pair[0] is None else align.query_qualities[pair[0]])
        if(ref_start_index is None):
            ref_start_index = pair[1]
        if(query_start_index is None):
            query_start_index = pair[0]

    if(ref_start_index is None):
        ref_start_index = 0
    if(query_start_index is None):
        query_start_index = 0
    return AlignStr("".join(ref_align_str_list), "".join(query_align_str_list), query_quality_list, ref_start_index, query_start_index)

def pairToSingle(align1:pysam.AlignedSegment, align2:pysam.AlignedSegment, ref_seq:str=None):
    '''
    没有重叠区则返回两个AlingStr类\n
    有重叠区但重叠区不一致则返回空值\n
    有重叠区且重叠区一致则返回合并后的AlignStr类
    '''
    align1_str = alignToStr(align1, ref_seq)
    align2_str = alignToStr(align2, ref_seq)
    if(align2_str.ref_start_index < align1_str.ref_start_index):
        align2_str,align1_str = [align1_str, align2_str]

    if(align2_str.ref_start_index > align1_str.ref_end_index):
        ############没有重叠区，返回列表
        return [align1_str, align2_str]
    else:
        ############有重叠区
        temp_align_str = align1_str.reference_sub(align2_str.ref_start_index,align1_str.ref_end_index)
        if(temp_align_str.ref_seq == align2_str.ref_seq[:len(temp_align_str)] and temp_align_str.query_seq == align2_str.query_seq[:len(temp_align_str)]):
            #########重叠区一致
            new_ref_seq = align1_str.ref_seq + align2_str.ref_seq[len(temp_align_str):]
            new_query_seq = align1_str.query_seq + align2_str.query_seq[len(temp_align_str):]
            return AlignStr(new_ref_seq, new_query_seq, align1_str.ref_start_index, align1_str.query_start_index)
        else:
            #########重叠区不一致
            return None


